import tensorflow as tf
from tensorflow.keras import layers, models
from tensorflow.keras.applications import ResNet50V2

class ModelBuilder:
    def __init__(self, input_shape=(224, 224, 3), num_classes=2):
        self.input_shape = input_shape
        self.num_classes = num_classes
        
    def build_basic_cnn(self):
        """构建基础CNN模型"""
        model = models.Sequential([
            # 第一个卷积块
            layers.Conv2D(32, (3, 3), activation='relu', input_shape=self.input_shape),
            layers.BatchNormalization(),
            layers.MaxPooling2D((2, 2)),
            
            # 第二个卷积块
            layers.Conv2D(64, (3, 3), activation='relu'),
            layers.BatchNormalization(),
            layers.MaxPooling2D((2, 2)),
            
            # 第三个卷积块
            layers.Conv2D(128, (3, 3), activation='relu'),
            layers.BatchNormalization(),
            layers.MaxPooling2D((2, 2)),
            
            # 全连接层
            layers.Flatten(),
            layers.Dropout(0.5),
            layers.Dense(256, activation='relu'),
            layers.BatchNormalization(),
            layers.Dense(self.num_classes, activation='softmax', kernel_regularizer=tf.keras.regularizers.l2(0.01))
        ])
        
        return model
        
    def build_transfer_model(self):
        """构建优化后的基于ResNet50V2的迁移学习模型"""
        # 定义输入层
        inputs = tf.keras.Input(shape=self.input_shape)
        
        # 添加数据增强层
        data_augmentation = tf.keras.Sequential([
            layers.RandomFlip("horizontal"),
            layers.RandomRotation(0.2),
            layers.RandomZoom(0.1),
            layers.RandomBrightness(0.2),
            layers.RandomContrast(0.2),
        ], name='data_augmentation')
        
        # 加载预训练模型
        base_model = ResNet50V2(
            weights='imagenet',
            include_top=False,
            input_shape=self.input_shape
        )
        base_model.trainable = False
        
        # 构建模型流程
        x = data_augmentation(inputs)
        x = base_model(x)
        x = layers.GlobalAveragePooling2D(name='pooling')(x)
        x = layers.Dense(128, 
                    activation='relu',
                    kernel_regularizer=tf.keras.regularizers.l2(0.01),
                    name='dense_1')(x)
        x = layers.BatchNormalization(name='batch_norm')(x)
        x = layers.Dropout(0.5, name='dropout')(x)
        
        # 二分类输出层
        outputs = layers.Dense(1, activation='sigmoid', name='output')(x)
        
        # 创建模型
        model = tf.keras.Model(inputs=inputs, outputs=outputs, name='object_detector')
        
        # 编译模型
        model.compile(
            optimizer=tf.keras.optimizers.Adam(learning_rate=5e-5),
            loss=tf.keras.losses.BinaryCrossentropy(label_smoothing=0.1),
            metrics=[
                'accuracy',
                tf.keras.metrics.Precision(),
                tf.keras.metrics.Recall()
            ]
        )
        
        # 打印模型信息
        print("\n模型结构:")
        print(f"- 输入形状: {self.input_shape}")
        print(f"- 任务类型: 二分类（目标物品 vs 非目标物品）")
        print(f"- 使用损失函数: BinaryCrossentropy")
        
        # 构建模型
        model.build(input_shape=(None,) + self.input_shape)
        model.summary()
        
        return model 